package model

import (
	"context"
	"fmt"
	"os"
	"strings"
	"time"

	"ai-rag/config"

	"github.com/cloudwego/eino-ext/components/model/openai"
	"github.com/cloudwego/eino/components/model"
	"github.com/cloudwego/eino/schema"
	"github.com/sirupsen/logrus"
)

// LLM 语言模型接口
type LLM interface {
	Generate(ctx context.Context, messages []*schema.Message) (*schema.Message, error)
	GenerateWithPrompt(ctx context.Context, systemPrompt, userQuery string) (*schema.Message, error)
}

// GenericLLM 通用语言模型实现
type GenericLLM struct {
	typeName     string
	modelPath    string
	temperature  float32
	maxTokens    int
	systemPrompt string
	einoModel    model.ChatModel
}

// NewGenericLLM 创建新的通用语言模型
func NewGenericLLM(modelType, modelPath string, temperature float32, maxTokens int, systemPrompt string) (*GenericLLM, error) {
	// 创建临时配置对象
	config := &config.ModelConfig{
		Type:         modelType,
		ModelPath:    modelPath,
		Temperature:  temperature,
		MaxTokens:    maxTokens,
		SystemPrompt: systemPrompt,
	}

	// 委托给配置驱动的初始化函数
	return NewGenericLLMFromConfig(config)
}

// NewGenericLLMFromConfig 从配置创建通用语言模型
func NewGenericLLMFromConfig(config *config.ModelConfig) (*GenericLLM, error) {
	logrus.WithField("model_type", config.Type).WithField("model_name", config.ModelName).Info("从配置初始化语言模型")

	// 创建聊天模型配置，完全从配置文件读取参数
	chatConfig := &openai.ChatModelConfig{
		MaxCompletionTokens: &config.MaxTokens,
		Temperature:         &config.Temperature,
	}

	// 设置超时时间
	timeout := 30 * time.Second // 默认超时
	if config.Timeout > 0 {
		timeout = time.Duration(config.Timeout) * time.Second
	}
	chatConfig.Timeout = timeout

	// 根据模型类型设置特定参数
	switch strings.ToLower(config.Type) {
	case "local":
		// 本地模型配置
		chatConfig.Model = getConfigValue(config.ModelPath, getEnvOrDefault("LOCAL_MODEL_PATH", ""))
		// 本地模型可能需要的API密钥
		apiKey := getEnvOrDefault("LOCAL_MODEL_API_KEY", "")
		if apiKey != "" {
			chatConfig.APIKey = apiKey
		}

	case "openai":
		// OpenAI模型配置
		chatConfig.APIKey = getEnvOrDefault("OPENAI_API_KEY", "")
		chatConfig.Model = getConfigValue(config.ModelName, getEnvOrDefault("OPENAI_MODEL", "gpt-4"))

		// 检查是否启用Azure OpenAI
		if config.AzureEnabled {
			chatConfig.ByAzure = true
			chatConfig.BaseURL = getConfigValue(config.AzureBaseURL, getEnvOrDefault("AZURE_OPENAI_ENDPOINT", ""))
			chatConfig.APIVersion = getConfigValue(config.AzureAPIVersion, getEnvOrDefault("AZURE_API_VERSION", "2024-02-15-preview"))
		}

	case "ollama":
		// Ollama模型配置
		chatConfig.Model = getConfigValue(config.ModelName, getEnvOrDefault("OLLAMA_MODEL", "llama2"))
		chatConfig.BaseURL = getConfigValue(config.BaseURL, getEnvOrDefault("OLLAMA_BASE_URL", "http://localhost:11434/api"))

	case "zhipuai":
		// 智谱AI模型配置
		chatConfig.APIKey = getEnvOrDefault("ZHIPU_API_KEY", "")
		chatConfig.Model = getConfigValue(config.ModelName, getEnvOrDefault("ZHIPU_MODEL", "glm-4"))
		// 设置智谱AI的基础URL如果有配置
		if config.BaseURL != "" {
			chatConfig.BaseURL = config.BaseURL
		}

	default:
		return nil, fmt.Errorf("不支持的模型类型: %s", config.Type)
	}

	// 日志记录最终使用的配置
	logrus.WithFields(logrus.Fields{
		"type":          config.Type,
		"model":         chatConfig.Model,
		"temperature":   config.Temperature,
		"max_tokens":    config.MaxTokens,
		"timeout":       timeout,
		"azure_enabled": chatConfig.ByAzure,
		"has_api_key":   chatConfig.APIKey != "",
	}).Debug("最终模型配置")

	ctx := context.Background()
	einoModel, err := openai.NewChatModel(ctx, chatConfig)
	if err != nil {
		logrus.WithError(err).Error("初始化模型失败")
		return nil, err
	}

	return &GenericLLM{
		typeName:     config.Type,
		modelPath:    config.ModelPath,
		temperature:  config.Temperature,
		maxTokens:    config.MaxTokens,
		systemPrompt: config.SystemPrompt,
		einoModel:    einoModel,
	}, nil
}

// getConfigValue 优先从配置获取值，如果为空则使用默认值
func getConfigValue(configValue, defaultValue string) string {
	if configValue != "" {
		return configValue
	}
	return defaultValue
}

// Generate 生成文本响应
func (m *GenericLLM) Generate(ctx context.Context, messages []*schema.Message) (*schema.Message, error) {
	logrus.WithField("message_count", len(messages)).Info("开始生成文本响应")

	response, err := m.einoModel.Generate(ctx, messages)
	if err != nil {
		logrus.WithError(err).Error("生成文本响应失败")
		return nil, err
	}

	logrus.WithField("response_length", len(response.Content)).Info("文本响应生成完成")
	return response, nil
}

// GenerateWithPrompt 使用系统提示和用户查询生成响应
func (m *GenericLLM) GenerateWithPrompt(ctx context.Context, systemPrompt, userQuery string) (*schema.Message, error) {
	// 如果没有提供系统提示，使用默认的
	if systemPrompt == "" {
		systemPrompt = m.systemPrompt
	}

	// 构建消息列表
	messages := []*schema.Message{
		schema.SystemMessage(systemPrompt),
		schema.UserMessage(userQuery),
	}

	// 生成响应
	return m.Generate(ctx, messages)
}

// GetEinoModel 获取底层的eino模型，用于多查询检索器等功能
func (m *GenericLLM) GetEinoModel() model.ChatModel {
	return m.einoModel
}

// getEnvOrDefault 获取环境变量，如果不存在则返回默认值
func getEnvOrDefault(key, defaultValue string) string {
	if value := os.Getenv(key); value != "" {
		return value
	}
	return defaultValue
}

// RAGLLM RAG专用语言模型，集成了检索和生成功能
type RAGLLM struct {
	llm          LLM
	systemPrompt string
}

// NewRAGLLM 创建新的RAG语言模型
func NewRAGLLM(llm LLM, systemPrompt string) *RAGLLM {
	return &RAGLLM{
		llm:          llm,
		systemPrompt: systemPrompt,
	}
}

// GenerateWithDocs 使用检索到的文档生成回答
func (rag *RAGLLM) GenerateWithDocs(ctx context.Context, query string, docs []*schema.Document) (*schema.Message, error) {
	logrus.WithField("query", query).WithField("docs_count", len(docs)).Info("开始基于文档生成回答")

	// 构建增强的系统提示
	enhancedSystemPrompt := rag.systemPrompt
	enhancedSystemPrompt += "\n\n检索到的相关文档信息：\n\n"

	// 添加检索到的文档
	for i, doc := range docs {
		enhancedSystemPrompt += fmt.Sprintf("文档 %d:\n%s\n\n", i+1, doc.Content)
	}

	// 生成回答
	response, err := rag.llm.GenerateWithPrompt(ctx, enhancedSystemPrompt, query)
	if err != nil {
		logrus.WithError(err).Error("基于文档生成回答失败")
		return nil, err
	}

	logrus.Info("基于文档的回答生成完成")
	return response, nil
}

// StreamingLLM 流式语言模型接口
type StreamingLLM interface {
	LLM
	GenerateStream(ctx context.Context, messages []*schema.Message, callback func(chunk string) error) error
}

// StreamingGenericLLM 流式通用语言模型实现
type StreamingGenericLLM struct {
	*GenericLLM
	streamingModel openai.ChatModel
}

// NewStreamingGenericLLM 创建新的流式通用语言模型
func NewStreamingGenericLLM(modelType, modelPath string, temperature float32, maxTokens int, systemPrompt string) (*StreamingGenericLLM, error) {
	genericLLM, err := NewGenericLLM(modelType, modelPath, temperature, maxTokens, systemPrompt)
	if err != nil {
		return nil, err
	}

	// 尝试转换为流式模型
	streamingModel, ok := genericLLM.einoModel.(*openai.ChatModel)
	if !ok {
		return nil, fmt.Errorf("模型 %s 不支持流式输出", modelType)
	}

	return &StreamingGenericLLM{
		GenericLLM:     genericLLM,
		streamingModel: *streamingModel,
	}, nil
}

// GenerateStream 流式生成文本响应
func (m *StreamingGenericLLM) GenerateStream(ctx context.Context, messages []*schema.Message, callback func(chunk string) error) error {
	logrus.WithField("message_count", len(messages)).Info("开始流式生成文本响应")

	// 定义内部回调，添加日志
	internalCallback := func(chunk string) error {
		logrus.WithField("chunk_length", len(chunk)).Debug("接收到流式输出块")
		return callback(chunk)
	}

	if err := m.streamingModel.GenerateStream(ctx, messages, internalCallback); err != nil {
		logrus.WithError(fmt.Errorf("%v", err)).Error("流式生成文本响应失败")
		return fmt.Errorf("%v", err)
	}

	logrus.Info("流式文本响应生成完成")
	return nil
}

// RAGStreamingLLM RAG专用流式语言模型
type RAGStreamingLLM struct {
	streamingLLM *StreamingGenericLLM
	systemPrompt string
}

// NewRAGStreamingLLM 创建新的RAG流式语言模型
func NewRAGStreamingLLM(streamingLLM *StreamingGenericLLM, systemPrompt string) *RAGStreamingLLM {
	return &RAGStreamingLLM{
		streamingLLM: streamingLLM,
		systemPrompt: systemPrompt,
	}
}

// GenerateStreamWithDocs 使用检索到的文档流式生成回答
func (rag *RAGStreamingLLM) GenerateStreamWithDocs(ctx context.Context, query string, docs []*schema.Document, callback func(chunk string) error) error {
	logrus.WithField("query", query).WithField("docs_count", len(docs)).Info("开始基于文档流式生成回答")

	// 构建增强的系统提示
	enhancedSystemPrompt := rag.systemPrompt
	enhancedSystemPrompt += "\n\n检索到的相关文档信息：\n\n"

	// 添加检索到的文档
	for i, doc := range docs {
		enhancedSystemPrompt += fmt.Sprintf("文档 %d:\n%s\n\n", i+1, doc.Content)
	}

	// 构建消息列表
	messages := []*schema.Message{
		schema.SystemMessage(enhancedSystemPrompt),
		schema.UserMessage(query),
	}

	// 流式生成回答
	if err := rag.streamingLLM.GenerateStream(ctx, messages, callback); err != nil {
		logrus.WithError(err).Error("基于文档流式生成回答失败")
		return err
	}

	logrus.Info("基于文档的流式回答生成完成")
	return nil
}
